// Package main 实现了一个WebSocket服务器，用于实时消息广播和客户端连接管理
package main

import (
	"context"  // 用于处理上下文和取消操作
	"fmt"      // 用于格式化输出
	"log"      // 用于日志记录
	"net/http" // 提供HTTP服务器和客户端功能
	"sync"     // 提供同步原语
	"time"     // 用于时间相关操作

	"github.com/gorilla/websocket" // WebSocket协议的实现
)

var (
	// clients 使用sync.Map存储所有活动的WebSocket连接
	// key: *websocket.Conn, value: bool
	clients sync.Map

	// broadcast 是用于向所有客户端广播消息的通道
	broadcast = make(chan []byte)

	// mutex 用于保护WebSocket连接的并发写操作
	mutex sync.Mutex
)

// upgrader 配置了WebSocket连接的升级器
var upgrader = websocket.Upgrader{
	// CheckOrigin 允许所有跨域请求
	// 在生产环境中应该根据需求限制来源
	CheckOrigin: func(r *http.Request) bool { return true },
}

const (
	// heartbeatInterval 定义了发送心跳消息的时间间隔
	heartbeatInterval = 10 * time.Second

	// pingInterval 定义了发送ping消息的时间间隔
	// 用于检测连接是否存活
	pingInterval = 1 * time.Minute
)

// main 函数是程序的入口点
// 初始化HTTP路由并启动WebSocket服务器
func main() {
	// 创建一个可取消的上下文，用于优雅关闭服务
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 设置路由处理函数
	http.HandleFunc("/", handleIndex)      // 处理首页请求
	http.HandleFunc("/ai/hcr", handleChat) // 处理WebSocket连接请求

	// 启动消息处理和ping检测的goroutine
	go handleMessages(ctx)
	go handlePing(ctx)

	// 启动HTTP服务器
	log.Println("Server started on localhost:8303")
	err := http.ListenAndServe(":8303", nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

// handleIndex 处理对根路径的HTTP GET请求
// 返回index.html页面
func handleIndex(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

// handleChat 处理WebSocket连接请求
// 负责：
// 1. 升级HTTP连接为WebSocket连接
// 2. 管理客户端连接的生命周期
// 3. 处理心跳消息
// 4. 接收客户端消息并广播
func handleChat(w http.ResponseWriter, r *http.Request) {
	// 将HTTP连接升级为WebSocket连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// 确保连接最终会被关闭，并从客户端列表中移除
	defer func() {
		conn.Close()
		clients.Delete(conn)
	}()

	// 将新连接添加到客户端列表
	clients.Store(conn, true)

	// 创建并启动心跳定时器
	heartbeatTimer := time.NewTimer(heartbeatInterval)
	defer heartbeatTimer.Stop()

	// 持续处理客户端消息和心跳
	for {
		select {
		case <-heartbeatTimer.C:
			// 发送心跳消息
			mutex.Lock()
			err := conn.WriteMessage(websocket.TextMessage, []byte("heartbeat"))
			mutex.Unlock()
			if err != nil {
				log.Println(err)
				return
			}
			// 重置心跳定时器
			heartbeatTimer.Reset(heartbeatInterval)

		default:
			// 读取客户端消息
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Println(err)
				return
			}
			fmt.Println("WebSocket服务 Received message:", string(message))
			// 将消息发送到广播通道
			broadcast <- message
		}
	}
}

// handleMessages 处理广播消息
// 从broadcast通道接收消息并发送给所有连接的客户端
// 参数:
//   - ctx: 用于控制goroutine生命周期的上下文
func handleMessages(ctx context.Context) {
	for {
		select {
		case message := <-broadcast:
			// 向所有客户端广播消息
			clients.Range(func(key, value interface{}) bool {
				client := key.(*websocket.Conn)
				mutex.Lock()
				err := client.WriteMessage(websocket.TextMessage, message)
				mutex.Unlock()
				if err != nil {
					log.Println(err)
					client.Close()
					clients.Delete(client)
				}
				return true
			})
		case <-ctx.Done():
			return
		}
	}
}

// handlePing 定期向所有客户端发送ping消息
// 用于检测连接是否存活
// 参数:
//   - ctx: 用于控制goroutine生命周期的上下文
func handlePing(ctx context.Context) {
	// 创建定时器
	ticker := time.NewTicker(pingInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// 向所有客户端发送ping消息
			clients.Range(func(key, value interface{}) bool {
				client := key.(*websocket.Conn)
				mutex.Lock()
				err := client.WriteMessage(websocket.PingMessage, nil)
				mutex.Unlock()
				if err != nil {
					log.Println(err)
					client.Close()
					clients.Delete(client)
				}
				return true
			})
		case <-ctx.Done():
			return
		}
	}
}
